#coding:utf8
from django.core.management.base import BaseCommand
import base64
import socket
import struct
import hashlib
import threading,random

connectionlist = {}

def deleteconnection(item):
    global connectionlist
    del connectionlist['connection'+item]
    
class WebSocket(threading.Thread):
    def __init__(self,conn,index,name,remote, path="/"):
        threading.Thread.__init__(self)
        self.conn = conn
        self.index = index
        self.name = name
        self.remote = remote
        self.path = path
        self.buffer = ""
        self.newprotocol = False
        
    def run(self):
        print 'Socket%s Start!' % self.index
        headers = {}
        self.handshaken = False

        while True:
            if self.handshaken == False:
                print 'Socket%s Start Handshaken with %s!' % (self.index,self.remote)
                self.buffer += self.conn.recv(1024)
                if self.buffer.find('\r\n\r\n') != -1:
                    header, data = self.buffer.split('\r\n\r\n', 1)
                    for line in header.split("\r\n")[1:]:
                        key, value = line.split(": ", 1)
                        headers[key] = value
                        
                    if headers.has_key("Sec-WebSocket-Key"):
                        print 'Socket%s, new version procotol, handshake with %s .' % (self.index,self.remote)
                        handshakestr = self.handshakeNew(headers,data)
                    else:
                        print 'Socket%s, old version procotol, handshake with %s .' % (self.index,self.remote)
                        handshakestr = self.handshakeOld(headers,data)
                    self.conn.send(handshakestr)
                    self.handshaken = True
                    print 'Socket%s Handshaken with %s success!' % (self.index,self.remote)
                    self.sendMessage('Welcome, '+self.name+' !')
            else:
                if self.newprotocol:
                    try:  
                        pData = self.conn.recv(8196)  
                        if not len(pData): 
                            continue
                    except:  
                        pass
                    else:  
                        code_length = ord(pData[1]) & 127  
                        if code_length == 126:  
                            masks = pData[4:8]  
                            data = pData[8:]  
                        elif code_length == 127:  
                            masks = pData[10:14]  
                            data = pData[14:]  
                        else:  
                            masks = pData[2:6]  
                            data = pData[6:]  
          
                        raw_str = ""  
                        i = 0  
                        for d in data:  
                            raw_str += chr(ord(d) ^ ord(masks[i%4]))  
                            i += 1
                            
                        self.sendMessage(raw_str)
                else:
                    self.buffer += self.conn.recv(64)
                    if self.buffer.find("\xFF")!=-1:
                        s = self.buffer.split("\xFF")[0][1:]
                        if s=='quit':
                            print 'Socket%s Logout!' % (self.index)
                            self.sendMessage(self.name+' Logout')
                            deleteconnection(str(self.index))
                            self.conn.close()
                            break
                        else:
                            print 'Socket%s Got msg:%s from %s!' % (self.index,s,self.remote)
                            self.sendMessage(self.name+':'+s)
                        self.buffer = ""
                    
    # handshakeNew 新版协议握手
    def handshakeNew(self,headers,data):
        szKey = base64.b64encode(hashlib.sha1(headers['Sec-WebSocket-Key'] + '258EAFA5-E914-47DA-95CA-C5AB0DC85B11').digest())  
        szHost = headers["Host"]  
        our_handshake = 'HTTP/1.1 101 Switching Protocols\r\n' \
                        'Upgrade:websocket\r\n'\
                        'Connection: Upgrade\r\n'\
                        'Sec-WebSocket-Accept:'+ szKey + '\r\n' \
                        'WebSocket-Location: ws://' + szHost + '/WebManagerSocket\r\n' \
                        'WebSocket-Protocol:WebManagerSocket\r\n\r\n'
        self.newprotocol = True
        return our_handshake
                    
    # handshakeOld 旧版协议握手
    def handshakeOld(self,headers,data):
        headers["Location"] = "ws://%s%s" %(headers["Host"], self.path)
        key1 = headers["Sec-WebSocket-Key1"]
        key2 = headers["Sec-WebSocket-Key2"]
        if len(data) < 8:
            data += self.conn.recv(8-len(data))
        key3 = data[:8]
        self.buffer = data[8:]
        token = self.generate_token(key1, key2, key3)
                     
        handshake = '\
                                HTTP/1.1 101 Web Socket Protocol Handshake\r\n\
                                Upgrade: WebSocket\r\n\
                                Connection: Upgrade\r\n\
                                Sec-WebSocket-Origin: %s\r\n\
                                Sec-WebSocket-Location: %s\r\n\r\n\
                                ' %(headers['Origin'], headers['Location'])
        return handshake + token
    
    def sendMessage(self,message):
        global connectionlist
        for connection in connectionlist.values():
            if self.newprotocol:
                connection.send(self.newversion_senddata(message))
            else:
                connection.send("\x00%s\xFF" % message)
            
    def newversion_senddata(self,data):
        back_str = []
        back_str.append('\x81')
        data_length = len(data)
        if data_length < 125:
            back_str.append(chr(data_length))
        else:
            back_str.append(chr(126))
            back_str.append(chr(data_length >> 8))
            back_str.append(chr(data_length & 0xFF))
        return "".join(back_str) + data
            
    def generate_token(self, key1, key2, key3):
        num1 = int("".join([digit for digit in list(key1) if digit.isdigit()]))
        spaces1 = len([char for char in list(key1) if char == " "])
        num2 = int("".join([digit for digit in list(key2) if digit.isdigit()]))
        spaces2 = len([char for char in list(key2) if char == " "])

        combined = struct.pack(">II", num1/spaces1, num2/spaces2) + key3
        return hashlib.md5(combined).digest()
    
        

class WebSocketServer(object):
    def __init__(self):
        self.socket = None
    def begin(self):
        print 'WebSocketServer Start!'
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.bind(("127.0.0.1",1234))
        self.socket.listen(50)
        
        global connectionlist
        
        i=0
        while True:
            connection, address = self.socket.accept()
            
            username=address[0]
            
            newSocket = WebSocket(connection,i,username,address)
            newSocket.start()
            connectionlist['connection'+str(i)]=connection
            i = i + 1

#TODO: function stop
class Command(BaseCommand):
    def handle(self, *args, **options):
        if len(args) == 1 and args[0] in ['start', 'stop', 'restart', 'run']:
            server = WebSocketServer()
            getattr(server, args[0])()
